package com.spark.sniffer;

import java.lang.instrument.Instrumentation;
import java.util.ArrayList;
import java.util.List;

import com.spark.sniffer.logging.api.Logger;
import com.spark.sniffer.logging.api.LoggerFactory;
import com.spark.sniffer.logging.core.EasyLogResolver;
import com.spark.sniffer.logging.schedule.LogFileCleaner;
import com.spark.sniffer.plugin.AbstractClassEnhancePluginDefine;
import com.spark.sniffer.plugin.PluginBootstrap;
import com.spark.sniffer.plugin.PluginFinder;
import net.bytebuddy.ByteBuddy;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.matcher.ElementMatchers;
import net.bytebuddy.utility.JavaModule;
import static net.bytebuddy.matcher.ElementMatchers.*;

/**
 * -javaagent:D:\workspace\xiao\project\backend-common\spark-sniffer-agent\target\spark-sniffer-agent.jar
 */
public class SparkSnifferAgent {
    private SparkSnifferAgent() {}

    private static final Logger log = LoggerFactory.getLogger(SparkSnifferAgent.class);

    public static void premain(String args, Instrumentation instrumentation) {
//         初始化启动参数
        AgentArgs.init(args);
        final PluginFinder pluginFinder;
        final PluginFinder bootstrapPluginFinder;
        List<AbstractClassEnhancePluginDefine> pluginDefines = new ArrayList<>();
        List<AbstractClassEnhancePluginDefine> bootstrapPluginDefines = new ArrayList<>();
        try {
            for(AbstractClassEnhancePluginDefine pluginDefine : new PluginBootstrap().loadPlugins()) {
                if(pluginDefine.isBootstrapInstrument()) {
                    bootstrapPluginDefines.add(pluginDefine);
                } else {
                    pluginDefines.add(pluginDefine);
                }
            }
            pluginFinder = new PluginFinder(pluginDefines);
            bootstrapPluginFinder = new PluginFinder(bootstrapPluginDefines);
        } catch (Exception e) {
            log.error(e, "Failed to load plugins: {}", e.getMessage());
            return;
        }
        if(!bootstrapPluginDefines.isEmpty()) {
            doInstanceInBootstrap(bootstrapPluginFinder, instrumentation);
        }
        if(!pluginDefines.isEmpty()) {
            doInstance(pluginFinder, instrumentation);
        }
        // start a sub thread to clean expired logs
        if(LoggerFactory.getResolver() instanceof EasyLogResolver) {
            new LogFileCleaner().start();
        }
    }

    private static void doInstanceInBootstrap(PluginFinder pluginFinder, Instrumentation instrumentation) {
        final ByteBuddy byteBuddy = new ByteBuddy();
        new AgentBuilder.Default(byteBuddy)
            .disableClassFormatChanges()
            .ignore(nameStartsWith("com.spark.sniffer")
                .or(nameStartsWith("com.intellij.rt.debugger.agent.CaptureStorage"))
                .or(nameStartsWith("net.bytebuddy."))
                .or(nameStartsWith("org.slf4j."))
                .or(nameStartsWith("org.apache.logging."))
                .or(nameStartsWith("org.groovy."))
                .or(nameContains("javassist"))
                .or(nameContains(".asm."))
                .or(nameStartsWith("sun.reflect"))
                .or(ElementMatchers.<TypeDescription>isSynthetic()))
            .type(pluginFinder.buildMatch())
            .transform(new Transformer(pluginFinder))
            .with(new Listener())
            .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
            .installOn(instrumentation);
    }

    private static void doInstance(PluginFinder pluginFinder, Instrumentation instrumentation) {
        final ByteBuddy byteBuddy = new ByteBuddy();
        new AgentBuilder.Default(byteBuddy)
            //disableClassFormatChanges()
            .ignore(nameStartsWith("com.spark.sniffer")
                .or(nameStartsWith("com.intellij.rt.debugger.agent.CaptureStorage"))
                .or(nameStartsWith("net.bytebuddy."))
                .or(nameStartsWith("org.slf4j."))
                .or(nameStartsWith("org.apache.logging."))
                .or(nameStartsWith("org.groovy."))
                .or(nameContains("javassist"))
                .or(nameContains(".asm."))
                .or(nameStartsWith("sun.reflect"))
                .or(ElementMatchers.<TypeDescription>isSynthetic()))
            .type(pluginFinder.buildMatch())
            .transform(new Transformer(pluginFinder))
            .with(new Listener())
            .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
            .installOn(instrumentation);
    }

    private static class Transformer implements AgentBuilder.Transformer {
        private PluginFinder pluginFinder;

        Transformer(PluginFinder pluginFinder) {
            this.pluginFinder = pluginFinder;
        }

        @Override
        public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription,
            ClassLoader classLoader, JavaModule module) {
            List<AbstractClassEnhancePluginDefine> pluginDefines = pluginFinder.find(typeDescription);

            if (pluginDefines.isEmpty()) {
                log.info("Matched class {}, but ignore by finding mechanism.", typeDescription.getTypeName());
                return builder;
            }

            DynamicType.Builder<?> newBuilder = builder;
            for (AbstractClassEnhancePluginDefine define : pluginDefines) {
                try {
                    DynamicType.Builder<?> possibleNewBuilder = define.define(typeDescription, newBuilder, classLoader);
                    if (possibleNewBuilder != null) {
                        newBuilder = possibleNewBuilder;
                    }
                } catch (Exception e) {
                    log.error(e, "Failed to transform {}: {}", typeDescription.getName(), e.getMessage());
                }
            }
            return newBuilder;
        }
    }

    private static class Listener implements AgentBuilder.Listener {
        public void onDiscovery(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {}

        public void onTransformation(TypeDescription typeDescription, ClassLoader classLoader, JavaModule module,
            boolean loaded, DynamicType dynamicType) {
            log.debug("Enhancing class {}...", typeDescription.getName());
        }

        public void onIgnored(TypeDescription typeDescription, ClassLoader classLoader, JavaModule module,
            boolean loaded) {}

        public void onError(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded,
            Throwable throwable) {
            log.error(throwable, "Error while enhancing class {}: {}", typeName, throwable.getMessage());
        }

        public void onComplete(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded) {}
    }

}
